from enum import Enum


class ClientStep(Enum):
    NONE = "step_none"  # 未开始
    LOGIN_DONE = "step_login_done"  # 登录完成
    INIT_DONE = "step_init_done"  # 初始化完成
    TRAIN_START = "step_train_start"  # 训练开始
    TRAIN_DONE = "step_train_done"  # 训练完成
    UPLOAD_DONE = "step_upload_done"  # 上传完成
    UPDATE_DONE = "step_update_done"  # 更新完成
    VERIFY_TRUE = "step_verify_true"  # 验证成功
    VERIFY_FALSE = "step_verify_false"  # 验证失败
    TRACE_DONE = "step_trace_done"  # 跟踪完成

    def get_description(self):
        """
        返回对应的中文描述
        """
        descriptions = {
            ClientStep.NONE: "未开始",
            ClientStep.LOGIN_DONE: "登录完成",
            ClientStep.INIT_DONE: "初始化完成",
            ClientStep.TRAIN_START: "训练开始",
            ClientStep.TRAIN_DONE: "训练完成",
            ClientStep.UPLOAD_DONE: "上传完成",
            ClientStep.UPDATE_DONE: "更新完成",
            ClientStep.VERIFY_TRUE: "验证成功",
            ClientStep.VERIFY_FALSE: "验证失败",
            ClientStep.TRACE_DONE: "跟踪完成",
        }
        return descriptions.get(self, "未知步骤")


class ClientData:
    _data = {}

    def __init__(self, id, user_id, user_name):
        self.id = id
        self.user_id = user_id
        self.user_name = user_name
        self.reset_status()

    def data(self):
        """
        获取客户端数据
        """
        if not self._data:
            self._data = {
                "user": {
                    "user_id": self.user_id,
                    "name": self.user_name,
                    "setting": "默认设置",
                    "dataset": "数据集一",
                    "client_id": self.id,
                }
            }
        return self._data

    def reset_status(self):
        """
        重置部分状态
        """
        self.set_user_data("agg_verify", True)
        self.set_user_data("status", "未训练")
        self.set_user_data("result1", "未验证")
        self.set_user_data("result2", "")
        self.set_user_data("setting", "默认设置")
        self.set_user_data("dataset", "数据集一")
        self.set_user_data("step_status", ClientStep.NONE.value)

    def set_step(self, step_name):
        self.set_user_data("step_status", step_name)

    def set_trace_step(self, step):
        self.data()["trace_step_status"] = step

    def set_agg_step(self, step):
        self.data()["agg_step_status"] = step

    def set_user_data(self, key, value):
        self.data()["user"][key] = value
